import numpy as np
from osgeo import gdal

def compute_diff_map(raster_file, n):
    """
    Compute the difference map, returning the difference between the maximum and minimum values of each pixel,
    excluding NoData areas.

    Parameters:
    raster_file (str): Input raster file path
    n (int): Window size radius, n=1 represents a 3x3 window

    Returns:
    diff_map (numpy.ndarray): Difference map (maximum value - minimum value for each pixel)
    """
    # Open raster data using GDAL
    dataset = gdal.Open(raster_file)
    if dataset is None:
        raise Exception("Failed to open raster file.")
    
    # Read raster data into numpy array
    band = dataset.GetRasterBand(1)  # Assuming single-band raster
    grid = band.ReadAsArray()
    
    # Get the dimensions of the grid
    rows, cols = grid.shape
    
    # Get NoData value
    nodata_value = band.GetNoDataValue()
    
    # Initialize the diff_map (same size as grid, filled with zeros)
    diff_map = np.zeros((rows, cols), dtype=np.float32)
    
    # Iterate over each pixel in the grid
    for i in range(rows):
        for j in range(cols):
            # Skip NoData areas
            if grid[i, j] == nodata_value:
                diff_map[i, j] = nodata_value
                continue
            
            # Define the boundaries of the search window (ensure within grid edges)
            row_min = max(0, i - n)
            row_max = min(rows - 1, i + n)
            col_min = max(0, j - n)
            col_max = min(cols - 1, j + n)
            
            # Collect all values in the search window
            search_unit = grid[row_min:row_max + 1, col_min:col_max + 1]
            
            # Exclude NoData values
            search_unit = search_unit[search_unit != nodata_value]
            
            # Skip if there are no valid values in the search window
            if len(search_unit) == 0:
                diff_map[i, j] = nodata_value
            else:
                # Compute the range (max - min), and assign it to the diff_map
                diff_map[i, j] = np.max(search_unit) - np.min(search_unit)
    
    return diff_map

def save_diff_map(output_file, diff_map, reference_raster):
    """
    Save the difference map as a raster file and remove NoData areas.

    Parameters:
    output_file (str): Output raster file path
    diff_map (numpy.ndarray): Difference map data
    reference_raster (gdal.Dataset): Reference raster dataset to get geographic information
    """
    # Get the geo-transform and projection from the reference raster
    geotransform = reference_raster.GetGeoTransform()
    projection = reference_raster.GetProjection()
    
    # Create a new raster dataset to store the diff_map
    driver = gdal.GetDriverByName('GTiff')  # Using GeoTIFF format
    if not driver:
        raise Exception("GDAL does not support GeoTIFF.")
    
    out_dataset = driver.Create(output_file, diff_map.shape[1], diff_map.shape[0], 1, gdal.GDT_Float32)
    out_dataset.SetGeoTransform(geotransform)
    out_dataset.SetProjection(projection)
    
    # Write the diff_map to the new raster band
    out_band = out_dataset.GetRasterBand(1)
    out_band.WriteArray(diff_map)
    
    # Set the NoData value
    out_band.SetNoDataValue(reference_raster.GetRasterBand(1).GetNoDataValue())
    
    # Close the dataset
    out_band.FlushCache()
    out_dataset.FlushCache()

# Example usage:
raster_file = 'Calculation_of_settlement_difference.tif'  # Replace with your file path
output_file = 'diff_map_output_n1.tif'  # Output file path for n=1
n = 1  # For a 3x3 window, n=1

# Compute the diff_map
dataset = gdal.Open(raster_file)
if dataset is None:
    raise Exception("Failed to open input raster file.")

diff_map_n1 = compute_diff_map(raster_file, n)

# Save the diff_map to a new raster file for n=1
save_diff_map(output_file, diff_map_n1, dataset)

# Print confirmation message
print(f"Difference map for n=1 saved as {output_file}")

# Repeat for n=2
output_file = 'diff_map_output_n2.tif'  # Output file path for n=2
n = 2  # For a 5x5 window, n=2

# Compute the diff_map for n=2
diff_map_n2 = compute_diff_map(raster_file, n)

# Save the diff_map to a new raster file for n=2
save_diff_map(output_file, diff_map_n2, dataset)

# Print confirmation message
print(f"Difference map for n=2 saved as {output_file}")

# Repeat for n=3
output_file = 'diff_map_output_n3.tif'  # Output file path for n=3
n = 3  # For a 7x7 window, n=3

# Compute the diff_map for n=3
diff_map_n3 = compute_diff_map(raster_file, n)

# Save the diff_map to a new raster file for n=3
save_diff_map(output_file, diff_map_n3, dataset)

# Print confirmation message
print(f"Difference map for n=3 saved as {output_file}")

